#ifndef __DIFFNORM_H
#define __DIFFNORM_H

#include <cassert>
#include <queue>
#include <sstream>
#include <unordered_map>
#include <unordered_set>
#include <vector>

#include "DiffNormAlgo.h"
#include "DiffNormUtils.h"

// Strategies for generating the candidates and sorting them.
enum CandidateGenerationStrategy {
    USAGE,  // estimated usage.
    GAIN    // estimated gain.
};

// Strategies for running DIFFNORM. 
enum DiffNormRunStrategy {
    RCPA,   // Regenerate the candidates post every candidate acceptance.
    RCTA    // Regenerate the candidates after testing all the candidates.
};

// Represents the meta information about each candidate tested.
struct CandidateStats {
    double estGain;
    double actGain;
    double size;         // total encoded size of the model and data after testing the candidate.
    uint32 estUsgCount;
    bool accepted;      // was this candidate accepted?
    string id;          // unique id of this candidate.

    CandidateStats() : estGain(0.0), actGain(0.0), size(0.0), 
                       estUsgCount(0), accepted(false), id("") {}

    CandidateStats(double estGain, double actGain, double size, 
                    uint32 estUsgCount, bool accepted, string id) :
                    estGain(estGain), actGain(actGain), size(size),
                    estUsgCount(estUsgCount), accepted(accepted), id(id) {}
};

// Wrapper around itemset. An itemset in a combined CT could belong to either 
// specific CT or common CT. If an itemset belongs to a specific CT CT_i
// , IsRecord will contain a pointer to the combined CT pertaining to CT_i i.e.,
// CT_i + CT_common. If an itemset belongs to a common CT, that itemset will be 
// in all the combined CTs and thus IsRecord will contain pointers to the 
// itemsets in all the combined CTs.
struct IsRecord {
    bool cleanupyourself;
    double costOfDeletion;   // how many bits would it cost to remove this record from CT?
    uint32 numCombinedCTs; 
    string dnid;            // unique id of the itemset in this record
    ItemSet **itemSetsInCombinedCTs;
    
    IsRecord(uint32 numCombinedCTs) : dnid(""), costOfDeletion(0.0), 
                        cleanupyourself(false), numCombinedCTs(numCombinedCTs) {
        itemSetsInCombinedCTs = new ItemSet*[numCombinedCTs];
        for (uint32 i = 0; i < numCombinedCTs; i++) 
            itemSetsInCombinedCTs[i] = NULL;
    }

    IsRecord(ItemSet* &itemSet, uint32 numCombinedCTs) : costOfDeletion(0.0), 
                        numCombinedCTs(numCombinedCTs), cleanupyourself(true) {
        dnid = itemSet->GetDNID();
        itemSetsInCombinedCTs = new ItemSet*[numCombinedCTs];
        for (uint32 i = 0; i < numCombinedCTs; i++)
            itemSetsInCombinedCTs[i] = itemSet->Clone();
    };

    ~IsRecord() {
        if(cleanupyourself) {
            for (uint32 i = 0; i < numCombinedCTs; ++i) {
                delete itemSetsInCombinedCTs[i];
            }
        }
        delete[] itemSetsInCombinedCTs;
        itemSetsInCombinedCTs = NULL;
    }

    inline void SetCostOfDeletion(double cost) {
        costOfDeletion = cost;
    }

    inline void SetItemSetAt(ItemSet *itemSet, uint32 combinedCTIdx) {
        dnid = itemSet->GetDNID();
        itemSetsInCombinedCTs[combinedCTIdx] = itemSet;
    }

    inline ItemSet* GetItemSetAt(uint32 combinedCTIdx) {
        return itemSetsInCombinedCTs[combinedCTIdx];
    }

    inline bool Equals(const IsRecord *r) {
        return this->dnid == r->dnid;
    }
};

typedef vector<IsRecord*>::iterator recordsItr;
typedef unordered_map<uint32, unordered_set<uint32>>::iterator uint32ToSetItr;
typedef pair<uint32, unordered_set<uint32>> uint32AndSetPair;

// Represents the candidate itemset for testing. A candidate itemset is created
// by merging the two itemsets. These two itemsets could be from the same CT or
// from different CTs. DNCandidate keeps the other meta information besides the
// merged itemset.
struct DNCandidate {
    uint32 leftCTIdx;
    uint32 rightCTIdx;
    uint32 estUsgCount;
    double estGain;
    double actGain;
    IsRecord *leftRecord;
    IsRecord *rightRecord;
    ItemSet *mergedIs;

    DNCandidate(uint32 leftCTIdx, uint32 rightCTIdx, IsRecord* &leftRecord,
            IsRecord* &rightRecord, uint32 numCombinedCTs, string mergedId,
            double estGain, uint32 estUsgCount, double actGain)
            : leftCTIdx(leftCTIdx), rightCTIdx(rightCTIdx), estUsgCount(0), 
            estGain(0.0), actGain(0.0), leftRecord(leftRecord), 
            rightRecord(rightRecord) { MergeRecords(numCombinedCTs, mergedId);}

    ~DNCandidate() {
        if (mergedIs != NULL) {
            delete mergedIs;
            mergedIs = NULL;
        }
    }

    void SetActualGain(double actGain) {
        actGain = actGain;
    }

    void MergeRecords(uint32 numCombinedCTs, const string &mergedId) {
        uint32 leftIsIdx = 0, rightIsIdx = 0;
        if (leftCTIdx != numCombinedCTs) leftIsIdx = leftCTIdx;
        if (rightCTIdx != numCombinedCTs) rightIsIdx = rightCTIdx;

        ItemSet *leftIs = leftRecord->GetItemSetAt(leftIsIdx);
        ItemSet *rightIs = rightRecord->GetItemSetAt(rightIsIdx);

        mergedIs = leftIs->Union(rightIs);
        mergedIs->SetDNID(mergedId);
    }

    inline string GetDNID() {
        return mergedIs->GetDNID();
    }

    inline uint32 GetLength() {
        return mergedIs->GetLength();
    }

    inline string GetLeftDNID() {
        return leftRecord->dnid;
    }

    inline string GetRightDNID() {
        return rightRecord->dnid;
    }

    inline uint32 GetSupport() {
        return mergedIs->GetSupport();
    }

    inline void SetEstimatedGain(double estimatedGain) {
        estGain = estimatedGain;
    }

    inline uint32 GetCardinality() {
        return mergedIs->GetLength();
    }

    bool operator==(const DNCandidate &c) {
        return this->mergedIs->Equals(c.mergedIs);
    }

    void PrettyPrint() {
        cout << "(" << leftRecord->dnid << "," << rightRecord->dnid << ")";
    }
};

// Represents the individual CTs (CT_is and CT_common). All the itemsets 
// are in combined code tables. PseudoCT keeps track of the itemsets pertaining
// to individual CTs.
struct PseudoCT {
    vector<IsRecord*> records;  // list of records
    IsRecord* deletedRecord;    // recently deleted record
    recordsItr deletedPos;      // position of the deleted record from the records container

    PseudoCT() {
        deletedRecord = NULL;
        deletedPos = records.end();
    };

    ~PseudoCT() {
        while (!records.empty()) {
            IsRecord *record = records.back();
            records.pop_back();
            delete record;
        }
    }

    uint32 NumRecords() {
        return records.size();
    }

    recordsItr Del(recordsItr &recordAt, bool zap) {
        if (zap) {
            delete *recordAt;
            *recordAt = NULL;
            deletedRecord = NULL;
            deletedPos = records.end();
        } else {
            deletedRecord = *recordAt;
            deletedPos = recordAt;
        }
        recordsItr nextItrPos = records.erase(recordAt);
        return nextItrPos;
    }

    void CommitLastDel() {
        delete deletedRecord;
        deletedRecord = NULL;
        deletedPos = records.end();
    }

    void RollbackLastDel() {
        records.insert(deletedPos, deletedRecord);
        deletedPos = records.end();
        deletedRecord = NULL;
    }

    void RollbackLastAdd() {
        IsRecord *lastRecord = records.back();
        records.pop_back();
        delete lastRecord;
        lastRecord = NULL;
    }

    void Add(IsRecord* &record) {
        records.push_back(record);
    }
};

// Main class for running DIFFNORM. All the important operational parts of 
// DIFFNORM are packed here.
class DiffNorm : public DiffNormAlgo {
public:
    DiffNorm(CodeTable **combinedCTs, uint32 numSplits, Config* config);
    ~DiffNorm();

    // Run the DiffNorm mining algorithm, main entry point
    void Run();

protected:
    uint32 mMinSup;
    uint32 mNumSingleCTs;
    uint32 mNumCombinedCTs;
    uint32 mNonSingletonCCTIdx;
    uint32 mSingletonCCTIdx; // CCT -> common code table
    uint32 mNumIterations;
    uint64 mNumGeneratedCandidates;
    uint64 mNumRejectedCandidates;

    bool mWriteReportFile;
    bool mWriteStatsFile;
    bool mReportTranslated;
    bool mAddCandidateToCCT;
    bool mCandidateAccepted;
    bool mGenerateMoreCandidates;
    bool mReEstimateAllCandidates;

    double mTotalSizeInStart;
    double mCompressionStartTime;
    double mCompressionEndTime;
    double mCurTotalSize;
    double mMinGain;

    CandidateGenerationStrategy mCandidateGenStrategy;
    DiffNormRunStrategy mRunStrategy;

    DNCandidate* mCandidate;
    vector<PseudoCT*> mSingleCTs; // For 2 split, CT_1, CT_2, CT_Omega(non-singleton), CT_Omega(singleton) 
    vector<DNCandidate*> mCandidates; // pool of candidates
    unordered_set<string> mGenerated; // recent generated candidates
    unordered_set<string> mRejected; // rejected candidates
    unordered_set<string> mAccepted; // accepted candidates
    vector<double> mCurrentCTSizes; // L(CT_1), L(CT_2), L(CT_Omega)
    vector<double> mCurrentSplitSizes; // L(D_1 | CT_1 + CT_Omega), L(D_2 | CT_2+CT_Omega)
    vector<double> mChangedCTSizes; // L(CT_1), L(CT_2), L(CT_Omega)
    vector<double> mChangedSplitSizes; // L(D_1 | CT_1 + CT_Omega), L(D_2 | CT_2+CT_Omega)
    vector<uint32> mCommitCTIndices; // commit to which code tables
    vector<CandidateStats> mCandidateStats; // statistics of all the candidates generated/tested

    inline bool AccceptedBefore(const string &dnid);
    void AcceptCandidate();
    void AddCandidateToAllCTs();
    inline bool AlreadyGenerated(const string &id);
    double CalcCompressionTime();
    double CalcCostOfDeletion(uint32 ctIdx, uint32 recordIdx);
    double CalcDbGain(DNCandidate *candidate, vector<uint32> &xyUsgList, vector<uint32> &xUsgList, vector<uint32> &yUsgList, vector<uint32> &xpUsgList, vector<uint32> &ypUsgList, uint32 combinedCtIdx);
    double CalcEncodedSizeOfCTFromSingleCT(uint32 idxOfSingleCT);
    double CalcEncodedSizeOfItemSetByST(ItemSet* itemSet);
    double CalcExpectation(ItemSet* itemSet);
    void CalcGainPostInsertion();
    double CalcGainFromIndividualSizes();
    void CalcImportanceOfItemSetsInAllCTs();
    double CalcTotalEncodedSize();
    double CalcTotalEncodedSizeFromIndividualSizes();
    void CheckGainInEachCTs();
    void CollectItemSetRecordsInCCT(vector<IsRecord*> &records);
    void CommitAcceptedChanges();
    void CommitChangesToCombinedCT(uint32 ctIdx);
    void CommitChangesToSingleCT(uint32 ctIdx);
    void CreateMergedId(const string &leftId, const string &rightId, string &mergedId);
    void EstimateGain(DNCandidate *candidate);
    uint32 EstimateMaxUsage(IsRecord *record1, IsRecord *record2, uint32 ctIdx1, uint32 ctIdx2);
    uint32 ExtractMinimumSupport(string iscName);
    void GatherUsagesForCommonCT(IsRecord *record, vector<uint32> &usages);
    inline void GenerateCandidates();
    inline void GenerateCandidateBasedOnUsage(IsRecord *leftRecord, IsRecord *rightRecord, ItemSet *leftIs, ItemSet *rightIs, uint32 leftCtIdx, uint32 rightCtIdx);
    void GenerateCandidatesBasedOnUsage();
    inline void GenerateCandidateBasedOnGain(IsRecord *leftRecord, IsRecord *rightRecord, ItemSet *leftIs, ItemSet *rightIs, uint32 leftCtIdx, uint32 rightCtIdx);
    void GenerateCandidatesBasedOnGain();
    CandidateGenerationStrategy GetCandidateGenerationStrategy(Config *config);
    DiffNormRunStrategy GetDiffNormRunStrategy(Config *config);
    inline uint32 GetItemSetIdx(uint32 ctIdx);
    DNCandidate* GetNextCandidate();
    inline uint32 GetNumAcceptedCandidates();
    inline uint32 GetNumRejectedCandidates();
    inline uint32 GetPrevUsageCount(ItemSet *is, uint32 ctIdx);
    uint32 GetPrevUsageCount(IsRecord *record, uint32 ctIdx);
    void GetRecords(uint32 ctIdx, vector<IsRecord*> &records);
    uint32 GetUsageCount(IsRecord *record, uint32 ctIdx);
    inline uint32 GetUsageCount(ItemSet *is, uint32 ctIdx);
    inline uint32* GetUsages(ItemSet *is, uint32 ctIdx);
    void GetUsageList(IsRecord *record, uint32 ctIdx, vector<uint32> &usgList);
    void GetUsageList(IsRecord *record, vector<uint32> &usgList, uint32 ctToAdd);
    void InitializeSingleCTs();
    uint8 IsZero(uint32 val);
    void LoadNextCandidate();
    inline void NoteDownCompressionStartTime();
    inline void NoteDownCompressionEndTime();
    bool PruneCommonCT();
    void PostPrune();
    bool PruneSpecificCT(uint32 ctIdx);
    inline bool RejectedBefore(ItemSet *is);
    inline bool RejectedBefore(const string &dnid);
    inline void RejectCandidate();
    void ReEstimateGainsOfCandidates();
    void ReEstimateGainsOfAllCandidates();
    void ReEstimateGainsOfChangedCandidates();
    void RemoveCandidateFromPool(const string &dnid);
    void RollbackRejectedChanges();
    void RunRegenerateCandidatesPostAcceptanceStrategy();
    void RunRegenerateCandidatesAfterTestingAllStrategy();
    inline void RunStrategy();
    void UpdateBookKeepingStuff();
    inline void UpdateCandidateStats();
    void UpdateSingletonsCCT();
    void WriteCandidateStatsToDisk();
    void WriteCTsToDisk();
    void WriteProgressToScreen();
    void WriteReportToDisk();
    void WriteReports();
};

#endif // __DIFFNORM_H
